import cv2
import os
import random
import numpy as np


# 输入文件夹路径
input_folder = "intestine/"
output_folder = "intestine-img/"

# 确保输出文件夹存在
os.makedirs(output_folder, exist_ok=True)

# 获取所有图片文件路径
image_files = [os.path.join(input_folder, f) for f in os.listdir(input_folder) if f.endswith(('.jpg', '.jpeg', '.png'))]
num_original_images = len(image_files)

# 需要生成的图像数量
target_num_images = 100
augment_per_image = target_num_images // num_original_images

# 定义数据增强函数
def augment_image(image):
    # 随机旋转
    angle = random.randint(-30, 30)
    M = cv2.getRotationMatrix2D((image.shape[1] // 2, image.shape[0] // 2), angle, 1)
    rotated = cv2.warpAffine(image, M, (image.shape[1], image.shape[0]))

    # 随机缩放
    scale = random.uniform(0.9, 1.1)
    scaled = cv2.resize(rotated, None, fx=scale, fy=scale)

    # 随机平移
    tx = random.randint(-10, 10)
    ty = random.randint(-10, 10)
    M = np.float32([[1, 0, tx], [0, 1, ty]])
    translated = cv2.warpAffine(scaled, M, (scaled.shape[1], scaled.shape[0]))

    # 随机水平翻转
    if random.choice([True, False]):
        translated = cv2.flip(translated, 1)

    return translated

# 开始数据扩增
image_counter = 0
for image_path in image_files:
    image = cv2.imread(image_path)
    base_name = os.path.basename(image_path).split('.')[0]

    for i in range(augment_per_image):
        augmented_image = augment_image(image)
        
        # 保存增强后的图像
        output_path = os.path.join(output_folder, f"{base_name}_aug_{i}.jpg")
        cv2.imwrite(output_path, augmented_image)
        image_counter += 1

        # 检查是否达到了目标图像数量
        if image_counter >= target_num_images:
            break

    if image_counter >= target_num_images:
        break

print(f"数据扩增完成，总共生成了 {image_counter} 张图像。")